#!/usr/bin/env python3
#
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import argparse
import sys
import shlex
import json
import socket
import time
import os
import copy

__version__='v1.1.9'

try:
    from shlex import quote
except ImportError:
    from pipes import quote

MODIFY_SUBCMD_DEVICE = 0
MODIFY_SUBCMD_QUEUE = 1

class JsonRpcSnapException(Exception):
    def __init__(self, message):
        self.message = message

class JsonRpcVirtnetClient(object):
    decoder = json.JSONDecoder()

    def __init__(self, address, port, timeout=60.0):
        self.sock = None
        self._request_id = 0
        self.timeout = timeout
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((address, int(port)))
        except socket.error as ex:
            raise JsonRpcSnapException("Error while connecting to %s:%s\n"
                                       "Error details: %s" % (address, port, ex))

    def __json_to_string(self, request):
        return json.dumps(request)

    def send(self, method, params=None):
        self._request_id += 1
        req = {
            'jsonrpc': '2.0',
            'method': method,
            'id': self._request_id
        }
        if params:
            req['params'] = copy.deepcopy(params)

        self.sock.sendall(self.__json_to_string(req).encode("utf-8"))
        return self._request_id

    def __string_to_json(self, request_str):
        try:
            obj, idx = self.decoder.raw_decode(request_str)
            return obj
        except ValueError:
            return None

    def recv(self):
        timeout = self.timeout
        start_time = time.process_time()
        response = None
        buf = ""

        while not response:
            try:
                timeout = timeout - (time.process_time() - start_time)
                self.sock.settimeout(timeout)
                buf += self.sock.recv(4096).decode("utf-8")
                response = self.__string_to_json(buf)
            except socket.timeout:
                break
            except ValueError:
                continue  # incomplete response; keep buffering

        self.sock.close()

        if not response:
            raise JsonRpcSnapException("Response Timeout")
        return response

    def call(self, method, params={}):
        if params:
            print(params)
        req_id = self.send(method, params)
        response = self.recv()

        if 'error' in response:
            params["method"] = method
            params["req_id"] = req_id
            msg = "\n".join(["request:", "%s" % json.dumps(params, indent=2),
                             "Got JSON-RPC error response",
                             "response:",
                             json.dumps(response['error'], indent=2)])
            raise JsonRpcSnapException(msg)

        return response['result']

def call_rpc_func(args):
    args.func(args)

def execute_script(parser, client, fd):
    executed_rpc = ""
    for rpc_call in map(str.rstrip, fd):
        if not rpc_call.strip():
            continue
        executed_rpc = "\n".join([executed_rpc, rpc_call])
        args = parser.parse_args(shlex.split(rpc_call))
        args.client = client
        try:
            call_rpc_func(args)
        except JsonRpcSnapException as ex:
            print("Exception:")
            print(executed_rpc.strip() + " <<<")
            print(ex.message)
            exit(1)

def hotplug(args):
    params = {
        'ib_device': args.ib_device,
    }
    if args.mac != None:
        params['mac'] = args.mac
    if args.mtu != None:
        params['mtu'] = args.mtu
    if args.num_queues != None:
        params['num_queues'] = args.num_queues
    if args.queue_size != None:
        params['queue_size'] = args.queue_size
    if args.features != None:
        params['features'] = args.features

    result = args.client.call('hotplug', params)
    print(json.dumps(result, indent=2))

def unplug(args):
    params = {
        'id': args.pf,
    }

    result = args.client.call('unplug', params)
    print(json.dumps(result, indent=2))

def list(args):
    result = args.client.call('list')
    print(json.dumps(result, indent=2))
    
def query(args):
    params = {}
    if args.all != None:
        params['all'] = args.all
    if args.vf != None:
        params['vf'] = args.vf
    if args.pf != None:
        params['pf'] = args.pf

    result = args.client.call('query', params)
    print(json.dumps(result, indent=2))

def modify(args):
    params = {
        'pf': args.pf,
    }
    if args.vf != None:
        params['vf'] = args.vf
    if args.subcmd == 'device':
        params['subcmd'] = MODIFY_SUBCMD_DEVICE
        if args.state != None:
            params['state'] = args.state
        if args.features != None:
            params['features'] = args.features
        if args.mac != None:
            params['mac'] = args.mac
        if args.mtu != None:
            params['mtu'] = args.mtu
        if args.link != None:
            params['link'] = args.link
        if args.rx_mode != None:
            params['rx_mode'] = args.rx_mode
    elif args.subcmd == 'queue':
        params['subcmd'] = MODIFY_SUBCMD_QUEUE
        if args.period_mode != None:
            params['period_mode'] = args.period_mode
        if args.period != None:
            params['period'] = args.period
        if args.max_count != None:
            params['max_count'] = args.max_count

    result = args.client.call('modify', params)
    print(json.dumps(result, indent=2))

def log_level(args):
    params = {
        'level': args.level,
    }

    result = args.client.call('log_level', params)
    print(json.dumps(result, indent=2))

def main():
    server_addr='127.0.0.1'
    server_port='12190'
    timeout=60.0

    parser = argparse.ArgumentParser(
        description='Nvidia virtio-net-controller command line interface ' + __version__)
    parser.add_argument('-v', '--version', action='version', version=__version__)
    subparsers = parser.add_subparsers(help='** Use -h for sub-command usage',
                                       dest='called_rpc_name')

    # Hotplug
    p = subparsers.add_parser('hotplug', help='hotplug virtnet device')
    p.add_argument('-i', '--ib_device', help="IB device name for creating "
                   "data path SF, e.g, mlx5_0",
                   required=True, type=str)
    p.add_argument('-m', '--mac', help="MAC in format 00:11:22:33:44:55",
                   required=True, type=str)
    p.add_argument('-t', '--mtu', help="MTU", required=True, type=int)
    p.add_argument('-n', '--num_queues', help="number of queues, at least 3",
                   required=True, type=int)
    p.add_argument('-s', '--queue_size', help="depth of each queue",
                   required=True, type=int)
    p.add_argument('-f', '--features', help="Feature bits in hex: "
                   "e.g, 0x80000 for vlan"
                   "\nVIRTIO_NET_F_CSUM             0"
                   "\nVIRTIO_NET_F_GSO              6"
                   "\nVIRTIO_NET_F_GUEST_TSO4       7"
                   "\nVIRTIO_NET_F_GUEST_TSO6       8"
                   "\nVIRTIO_NET_F_HOST_TSO4        11"
                   "\nVIRTIO_NET_F_HOST_TSO6        12"
                   "\nVIRTIO_NET_F_CTRL_VLAN        19"
                   "\nVIRTIO_NET_F_GUEST_ANNOUNCE   21",
                   required=False, type=str)
    p.set_defaults(func=hotplug)

    # Unplug
    p = subparsers.add_parser('unplug', help='unplug virtnet device')
    p.add_argument('-p', '--pf', help="PF device ID", required=True, type=int)
    p.set_defaults(func=unplug)

    # List
    p = subparsers.add_parser('list', help='list all virtnet devices')
    p.set_defaults(func=list)
    
    # Query
    p = subparsers.add_parser('query', help='query all or individual virtnet '
                              'device(s)')
    p.add_argument('-a', '--all', help='query all virtnet devices, including '
                    'both PFs and VFs', required=False, action="store_const",
                    const=1, default=0)
    p.add_argument('-p', '--pf', help="PF ID, refer to list command", 
                   required=False, type=int)
    p.add_argument('-v', '--vf', help="VF ID, refer to list command", 
                   required=False, type=int)
    p.set_defaults(func=query)

    # Modify
    p = subparsers.add_parser('modify', help='modify virtnet device')
    p.add_argument('-p', '--pf', help="pf ID", required=True, type=int)
    p.add_argument('-v', '--vf', help="vf ID", required=False, type=int)
    sp = p.add_subparsers()
    dp = sp.add_parser('device', help='modify device related options');
    dp.add_argument('-m', '--mac', help="MAC in format 00:11:22:33:44:55",
                   required=False, type=str)
    dp.add_argument('-t', '--mtu', help="MTU", required=False, type=int)
    dp.add_argument('-l', '--link', help="link status, 0: down, 1: up",
                   required=False, type=int)
    dp.add_argument('-s', '--state', help="device state",
                   required=False, type=int)
    dp.add_argument('-f', '--features', help="Feature bits in hex: "
                   "e.g, 0x80000 for vlan"
                   "\nVIRTIO_NET_F_CSUM             0"
                   "\nVIRTIO_NET_F_GSO              6"
                   "\nVIRTIO_NET_F_GUEST_TSO4       7"
                   "\nVIRTIO_NET_F_GUEST_TSO6       8"
                   "\nVIRTIO_NET_F_HOST_TSO4        11"
                   "\nVIRTIO_NET_F_HOST_TSO6        12"
                   "\nVIRTIO_NET_F_CTRL_VLAN        19"
                   "\nVIRTIO_NET_F_GUEST_ANNOUNCE   21",
                   required=False, type=str)
    dp.add_argument('-r', '--rx_mode', help="mode bits in hex, e.g. 0x1 "
                   "0:promisc 1:all-multi 2:all-uni 3:no-multi "
                   "4:no-uni 5:no-cast",
                   required=False, type=str)
    qp = sp.add_parser('queue', help='modify queue related options');
    qp.add_argument('-e', '--period_mode', help="period mode: event or cqe",
                   required=True, type=str, choices=['event', 'cqe'])
    qp.add_argument('-n', '--period', help="period (usec)",
                   required=True, type=int)
    qp.add_argument('-c', '--max_count', help="max count",
                   required=True, type=int)
    p.set_defaults(func=modify)
    dp.set_defaults(func=modify, subcmd='device')
    qp.set_defaults(func=modify, subcmd='queue')

    # Log
    p = subparsers.add_parser('log', help='set log level')
    p.add_argument('-l', '--level', help="log level: info/err/debug",
                    required=True, type=str, choices=['info', 'err', 'debug'])
    p.set_defaults(func=log_level)

    args = parser.parse_args()
    args.client = JsonRpcVirtnetClient(server_addr, server_port, timeout)
    if hasattr(args, 'func'):
        try:
            call_rpc_func(args)
        except JsonRpcSnapException as ex:
            print(ex)
            exit(1)
    elif sys.stdin.isatty():
        # No arguments and no data piped through stdin
        parser.print_help()
        exit(1)
    else:
        execute_script(parser, args.client, sys.stdin)

if __name__ == "__main__":
    main()
